import torch
import torch.nn as nn
import codecs

CHECKPOINT_PATH = "./data/seq2seq_ckpt.pth"

# 模型参数。必须与训练时的模型参数保持一致。
# LSTM的隐藏层规模
HIDDEN_SIZE = 1024
# 深层循环神经网络中LSTM结构的层数
NUM_LAYERS = 2
# 源语言词汇表大小
SRC_VOCAB_SIZE = 10000
# 目标语言词汇表大小
TRG_VOCAB_SIZE = 4000
# 在Softmax层和词向量层之间共享参数
SHARE_EMB_AND_SOFTMAX = True

# 词汇表中<sos>和<eos>的ID。在解码过程中需要用<sos>作为第一步的输入，并将检查是否是<eos>，因此需要知道这两个符号的ID
SOS_ID = 1
EOS_ID = 2

# 词汇表文件
SRC_VOCAB = "./data/en.vocab"
TRG_VOCAB = "./data/zh.vocab"


class NMTModel(nn.Module):
    def __init__(self):
        super(NMTModel, self).__init__()
        # 定义编码器和解码器所使用的LSTM结构
        self.encoder = nn.LSTM(HIDDEN_SIZE, HIDDEN_SIZE, NUM_LAYERS, batch_first=True)
        self.decoder = nn.LSTM(HIDDEN_SIZE, HIDDEN_SIZE, NUM_LAYERS, batch_first=True)

        # 为源语言和目标语言分别定义词向量
        self.src_embedding = nn.Embedding(SRC_VOCAB_SIZE, HIDDEN_SIZE)
        self.trg_embedding = nn.Embedding(TRG_VOCAB_SIZE, HIDDEN_SIZE)

        # 定义softmax层的变量
        if SHARE_EMB_AND_SOFTMAX:
            self.fc_out = nn.Linear(HIDDEN_SIZE, TRG_VOCAB_SIZE)
            self.fc_out.weight = self.trg_embedding.weight
        else:
            self.fc_out = nn.Linear(HIDDEN_SIZE, TRG_VOCAB_SIZE)

    def forward(self, src_input, max_dec_len=100):
        # 虽然输入只有一个句子，但因为dynamic_rnn要求输入是batch的形式，因此这里将输入句子整理为大小为1的batch
        src_size = torch.tensor([len(src_input)], dtype=torch.long)
        src_input = torch.tensor([src_input], dtype=torch.long)
        src_emb = self.src_embedding(src_input)

        # 使用dynamic_rnn构造编码器。这一步与训练时相同
        enc_outputs, enc_state = self.encoder(src_emb)

        # 设置解码的最大步数。这是为了避免在极端情况出现无限循环的问题。
        trg_ids = [SOS_ID]
        dec_state = enc_state

        for step in range(max_dec_len - 1):
            # 读取最后一步输出的单词，并读取其词向量
            trg_input = torch.tensor([[trg_ids[-1]]], dtype=torch.long)
            trg_emb = self.trg_embedding(trg_input)
            # 这里不使用dynamic_rnn，而是直接调用dec_cell向前计算一步。
            dec_output, dec_state = self.decoder(trg_emb, dec_state)
            # 计算每个可能的输出单词对应的logit，并选取logit值最大的单词作为这一步的输出。
            output = dec_output.squeeze(0)
            logits = self.fc_out(output)
            next_id = torch.argmax(logits).item()
            # 将这一步输出的单词写入循环状态的trg_ids中
            trg_ids.append(next_id)
            if next_id == EOS_ID:
                break

        return trg_ids


def main():
    # 定义训练用的循环神经网络模型
    model = NMTModel()

    # 加载预训练模型权重
    checkpoint = torch.load(CHECKPOINT_PATH, map_location='cpu', weights_only=True)
    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
    model.eval()

    # 定义一个测试的例子
    test_sentence = "This is a test ."
    print(test_sentence)

    # 根据英文词汇表，将测试句子转为单词ID。结尾加上<eos>的编号
    test_sentence = test_sentence + " <eos>"

    with codecs.open(SRC_VOCAB, 'r', 'utf-8') as vocab:
        src_vocab = [w.strip() for w in vocab.readlines()]
        # 运用dict，将单词和id对应起来组成字典，用于后面的转换
        src_id_dict = dict((src_vocab[x], x) for x in range(SRC_VOCAB_SIZE))
    test_en_ids = [(src_id_dict[en_text] if en_text in src_id_dict else src_id_dict['<unk>'])
                   for en_text in test_sentence.split()]
    print(test_en_ids)

    # 建立解码所需的计算图
    output_ids = model.forward(test_en_ids)

    # 根据中文词汇表，将翻译结果转换为中文文字。
    with codecs.open(TRG_VOCAB, "r", "utf-8") as f_vocab:
        trg_vocab = [w.strip() for w in f_vocab.readlines()]
    output_text = ''.join([trg_vocab[x] for x in output_ids[1:-1]])

    # 输出翻译结果
    print(output_text)


if __name__ == "__main__":
    main()



